import { io, Socket } from 'socket.io-client';
import { ref } from 'vue';

class WebSocketService {
    private socket: Socket | null = null;
    private connected = ref(false);
    private listeners = new Map<string, Function[]>();

    connect() {
        if (this.socket?.connected) {
            console.log('WebSocket already connected');
            return;
        }

        console.log('Connecting to WebSocket...');

        this.socket = io('/generation', {
            path: '/socket.io',
            transports: ['websocket', 'polling'],
            reconnection: true,
            reconnectionAttempts: 5,
            reconnectionDelay: 1000
        });

        this.setupEventListeners();
    }

    private setupEventListeners() {
        if (!this.socket) return;

        this.socket.on('connect', () => {
            console.log('WebSocket connected, socket id:', this.socket?.id);
            this.connected.value = true;
        });

        this.socket.on('disconnect', () => {
            console.log('WebSocket disconnected');
            this.connected.value = false;
        });

        this.socket.on('error', (error) => {
            console.error('WebSocket error:', error);
        });

        // 设置生成相关的事件监听
        ['generation_progress', 'generation_complete', 'generation_error'].forEach(event => {
            this.socket?.on(event, (data) => {
                this.notifyListeners(event, data);
            });
        });
    }

    subscribeToTask(promptId: string) {
        if (!this.socket?.connected) {
            console.warn('WebSocket not connected, cannot subscribe to task');
            return;
        }

        console.log('Subscribing to task:', promptId);
        this.socket.emit('subscribe_task', promptId);
    }

    on(event: string, callback: Function) {
        if (!this.listeners.has(event)) {
            this.listeners.set(event, []);
        }
        console.log('on event:', event, callback);
        this.listeners.get(event)?.push(callback);
    }

    off(event: string, callback: Function) {
        const callbacks = this.listeners.get(event);
        if (callbacks) {
            const index = callbacks.indexOf(callback);
            if (index !== -1) {
                callbacks.splice(index, 1);
            }
        }
    }

    private notifyListeners(event: string, data: any) {
        const callbacks = this.listeners.get(event);
        if (callbacks) {
            callbacks.forEach(callback => callback(data));
        }
    }

    disconnect() {
        if (this.socket) {
            this.socket.disconnect();
            this.socket = null;
            this.connected.value = false;
        }
    }

    isConnected() {
        return this.connected.value;
    }
}

export const wsService = new WebSocketService(); 